import json
import os
import time
import rpyc
import socket
from Env import action
from Env.obs import Observation
from functools import wraps
RUN_WINDOWS = True

def timecnt(a_func):
    @wraps(a_func)
    def wrapTheFunction(*args, **kwargs):
        ret = a_func(*args, **kwargs)
        return ret

    return wrapTheFunction


class EnvMgr:
    def __init__(self, env_id, ports, image_name, scen_name):
        self.ports = [3641, 5454]
        self.image_name = image_name
        self.scen_name = scen_name
        self.docker_name = f'combat-{env_id}'
        self.host = '127.0.0.1'

    def start_docker(self, volume_list=[]):
        docker_run = f'docker run -itd --hostname trsuser5 -p 6100:3641 -p 6111:5454 -p 6122:5455 -p 6113:5901 -p 6114:3328 -p 6114:3328/udp --name {self.docker_name} {self.image_name} /root/startup.sh '
        print(f'\n{docker_run}')
        os.system(docker_run)
        time.sleep(10)
        self.host = os.popen("docker inspect --format '{{ .NetworkSettings.IPAddress }}' " + self.docker_name).read().strip()
        self.__open()
        time.sleep(2)
        self.__start()
        time.sleep(5)

    def stop_docker(self):
        docker_stop = f'docker stop {self.docker_name}'
        os.system(docker_stop)
        docker_rm = f'docker rm -f {self.docker_name}'
        os.system(docker_rm)

    def restart_docker(self):
        docker_rs = f'docker restart {self.docker_name}'
        os.system(docker_rs)
        time.sleep(10)
        self.host = os.popen("docker inspect --format '{{ .NetworkSettings.IPAddress }}' " + self.docker_name).read().strip()
        self.__open()
        time.sleep(2)
        self.__start()
        time.sleep(5)

    def sock_send(self, ctrl_cmd):
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        print(ctrl_cmd, self.host)
        self.host = 'localhost'
        self.port = 6100
        sock.connect((self.host, self.port))#3641))#self.ports[0]))
        sock.send(bytes(ctrl_cmd, encoding="utf-8"))
        #res = sock.recv(256)
        #print(res.decode())
        sock.close()

    def __close(self):
        print('close')
        self.sock_send('CLOSE\n')

    def __open(self):
        print('open')
        self.sock_send(f"OPEN*{self.scen_name}*server\n")

    def __start(self):
        print('start')
        self.sock_send("START\n")

    def __stop(self):
        print('stop')
        self.sock_send("STOP\n")

    def _pause(self):
        self.sock_send("PAUSE\n")

    def _resume(self):
        self.sock_send("RESUME\n")



class EnvClient:
    def __init__(self, env_id, base_port, image_name, scen_name):
        if RUN_WINDOWS:
            return
        self.ports = [base_port+env_id*5+i for i in range(2)]
        self.env_mgr = EnvMgr(env_id, self.ports, image_name, scen_name)

    def reset(self):
        self.cmd_list = []
        self.lost_units = []
        if RUN_WINDOWS:
            import os
            #host = os.popen("cat /etc/resolv.conf |grep nameserver |awk '{print $2}'").read().strip()
            host = 'localhost'
            port = 5454#6101
            self.conn = rpyc.connect(host, port)
            return True

        print('Reset : restart the docker')
        self.env_mgr.stop_docker()
        self.env_mgr.start_docker()
        self.env_mgr.restart_docker()
        self.host = self.env_mgr.host   #127.0.0.1'
        self.port = 5454                #self.ports[1]
        start = time.time()
        while True:
            try:
                print('Try connect!')
                self.host = 'localhost'
                self.port = 6111
                self.conn = rpyc.connect(self.host, self.port)
                break
            except Exception as e:
                print(f'rpyc connect failed : {e} [{self.host}:{self.port}]')
                time.sleep(5)
                if time.time() - start > 100:
                    return False
        #self.__pause()
        print('Connect Success')
        return True

    def get_observation(self):
        response = self.conn.root.get_state()
        #print(help(type(response)))
        # print("Get Obs: type is :", type(response))
        ob_data = json.loads(response['json_data'])
        obs = Observation(ob_data)
        #obs.print_state()
        for unit in obs.unit_dict.values():
            if unit.isalive != 1 and unit.uid not in self.lost_units:
                self.lost_units.append(unit.uid)
                print(f'{unit.name}-{unit.uid} is lost')
        return obs

    def add_action(self, cmd):
        self.cmd_list.append(cmd)

    def take_action(self, cmd_list=None):
        if cmd_list == None:
            cmd_list = self.cmd_list
        data_json = json.dumps(cmd_list)
        #print(data_json)
        try:
            self.conn.root.take_action(data_json)
        except Exception as e:
            print("Exception", e)
        self.cmd_list = []

    def speedup(self, speed):
        cmd = action.make_simulation("SPEED", speed_num=speed)
        self.take_action([cmd])

    def __pause(self):
        # print("Resume : speed up 1")
        # self.speedup(1)
        if RUN_WINDOWS:
            cmd = action.make_simulation("PAUSE")
            self.take_action([cmd])
            return
        self.env_mgr._pause()

    def __resume(self):
        # time.sleep(10)
        # print("Resume : speed up 0")
        # self.speedup(0)
        if RUN_WINDOWS:
            cmd = action.make_simulation("RESUME")
            self.take_action([cmd])
            return
        self.env_mgr._resume()

    @timecnt
    def _get_current_time(self):
        tm = self.conn.root.get_time()
        # print('get current time:', type(tm))
        return int(tm['time'])

    def run_step(self, target_time):
        if len(self.cmd_list):
            self.take_action(self.cmd_list)
        self.__resume()
        try_times = 0
        while True:
            try:
                ctime = self._get_current_time()
                # print('time pass?', ctime, target_time)
                if ctime > target_time : #- now > diff:
                    break
            except Exception as e:
                try_times += 1
                print(f'Get Time Error : {e}')
                if try_times > 100:
                    raise Exception('Get times failed for too many times')
            time.sleep(0.01)
        self.__pause()
        return ctime



def run(rpyc_port):
    env_client = None
    try:
        if env_client is None:
            env_client = EnvClient(0, 6100, 'registry.inspir.ai:5000/lzy-combat:2.0', '/home/6.7(2).ntedt')

        for i in range(3):
            env_client.reset()
            env_client.speedup(0)
            print('Env Client Reset!')
            tm = 0
            for i in range(200):
                tm = env_client._get_current_time()
                env_client.run_step(i*20)
                print('Sim Time: ', tm)

            #print(json.dumps(state, ensure_ascii=False, indent=2))
    except Exception as e:
        env_client = None
        print("rpyc connect except: {}".format(e))


if __name__ == '__main__':
    time.sleep(1)
    run(5454)

